'''
Created on Jul 5, 2012

@author: nzhao
'''
import numpy as np

from traits.api import HasTraits, Float, Array

from math_qm.Constant import X, Y, Z

mix_state_2_spin = 0.25 * np.identity(4)
flip2 = np.kron(Y, Y)

def rootM(mat):
    w, m = np.linalg.eig(mat)
    sq = np.diag( np.sqrt( (1.0+0.0j) * w) )
    res = np.dot(m, np.dot(sq, np.conjugate(m).T ) )
    return res
    

class DensityMatrix(HasTraits):
    density_matrix = Array()
    
class DensityMatrix_1Qubit(DensityMatrix, HasTraits):
    state_I = Float(1.0)
    state_X = Float(0.0)
    state_Y = Float(0.0)
    state_Z = Float(0.0)
    
    def __init__(self):
        self.density_matrix = np.array([[0.5, 0.0], [0.0, 0.5]])
        self.calculate_vector()
        
    def calculate_vector(self):
        self.state_I = 1.0
        self.state_X = 0.5 * np.trace( np.dot(self.density_matrix, X) ).real
        self.state_Y = 0.5 * np.trace( np.dot(self.density_matrix, Y) ).real
        self.state_Z = 0.5 * np.trace( np.dot(self.density_matrix, Z) ).real
    def calculate_density_mat(self):
        self.density_matrix = 0.5 * np.identity(2) + self.state_X * X + self.state_Y * Y + self.state_Z * Z
        
    def set_vector(self, p0, px, py, pz):
        if self.state_I != 1.0:
            print 'error'
            self.state_I = 1.0
            
        self.state_X, self.state_Y, self.state_Z, =  px, py, pz
        self.calculate_density_mat()

    def get_vector(self):
        return [self.state_I, self.state_X, self.state_Y, self.state_Z]
    
    def set_density_mat(self, mat):
        self.density_matrix = mat
        self.calculate_vector()
        
    def get_density_mat(self):
        return self.density_matrix
    
    def _density_matrix_changed(self):
        self.calculate_vector()
    def _state_I_changed(self):
        self.calculate_density_mat()
    def _state_X_changed(self):
        self.calculate_density_mat()
    def _state_Y_changed(self):
        self.calculate_density_mat()
    def _state_Z_changed(self):
        self.calculate_density_mat()
        
class DensityMatrix_2Qubit(DensityMatrix, HasTraits):
    def __init__(self, density_matrix = mix_state_2_spin):
        self.density_matrix = density_matrix
    def concurrence(self):
        rho_1 = np.dot(flip2, np.dot(self.density_matrix.conj(), flip2) )
        root_rho = rootM(self.density_matrix)
        prod =  np.dot(root_rho, np.dot(rho_1, root_rho))

        Rmat = rootM(prod)
        w = np.sort( np.linalg.eigvals(Rmat) )

        res = max(0., w[3] - w[0] - w[1] -w[2])
        return res